package org.acghub.mtdb.core.engine;

import com.mongodb.client.*;
import com.mongodb.client.model.ReplaceOptions;
import lombok.extern.slf4j.Slf4j;
import org.acghub.mtdb.common.Config;
import org.acghub.mtdb.common.Table;
import org.acghub.mtdb.common.exception.LifecycleException;
import org.acghub.mtdb.core.StorageEngine;
import org.acghub.mtdb.core.StorageEngineType;
import org.bson.Document;

import java.util.Objects;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * @author zhouhe
 */

@Slf4j
public class DocumentMongoDBStorageEngine implements StorageEngine<Document, Document> {

    public static StorageEngine INSTANCE = new DocumentMongoDBStorageEngine();

    public static final String DEFAULT_CONNECTION = "mongodb://localhost:27017";
    public static final String DEFAULT_DATABASE = "mtdb_test";
    private final AtomicBoolean inited = new AtomicBoolean(false);
    private final AtomicBoolean closed = new AtomicBoolean(false);
    private MongoClient mongoClient;
    private MongoDatabase mongoDatabase;

    private MongoCollection<Document> getCollection(String name) {
        return mongoDatabase.getCollection(name);
    }

    private Document addId(Document key, Document value) {
        return value.append("_id", key.get("_id"));
    }

    @Override
    public void init(Config config) throws LifecycleException {
        try {
            if (inited.compareAndSet(false, true)) {
                String connectionStr = config.getStorageEngineConfig().getStr("connection", DEFAULT_CONNECTION);
                String database = config.getStorageEngineConfig().getStr("database", DEFAULT_DATABASE);
                mongoClient = MongoClients.create(connectionStr);
                mongoDatabase = mongoClient.getDatabase(database);
            }
        } catch (Exception e) {
            log.error("DocumentMongoDBStorageEngine init error, ", e);
            throw new LifecycleException("DocumentMongoDBStorageEngine init error, ", e);
        }
    }

    @Override
    public void close() throws LifecycleException {
        try {
            if (closed.compareAndSet(false, true)) {
                if (mongoDatabase != null) {
                    mongoDatabase = null;
                }
                if (mongoClient != null) {
                    mongoClient.close();
                    mongoClient = null;
                }
            }
        } catch (Exception e) {
            log.error("DocumentMongoDBStorageEngine close error, ", e);
            throw new LifecycleException(e);
        }
    }

    @Override
    public StorageEngineType type() {
        return StorageEngineType.DOCUMENT_MONGODB;
    }

    @Override
    public boolean exists(Table table, Document key) {
        FindIterable<Document> documents = getCollection(table.getName()).find(key).limit(1);
        return Objects.nonNull(documents.first());
    }

    @Override
    public Document select(Table table, Document key) {
        FindIterable<Document> documents = getCollection(table.getName()).find(key).limit(1);
        return documents.first();
    }

    @Override
    public boolean insert(Table table, Document key, Document value) {
        if (exists(table, key)) {
            return false;
        }
        getCollection(table.getName()).insertOne(addId(key, value));
        return true;
    }

    @Override
    public boolean update(Table table, Document key, Document value) {
        return Objects.nonNull(getCollection(table.getName()).findOneAndUpdate(key, value));
    }

    @Override
    public void upsert(Table table, Document key, Document value) {
        getCollection(table.getName()).replaceOne(key, addId(key, value), new ReplaceOptions().upsert(true));
    }

    @Override
    public void delete(Table table, Document key) {
        getCollection(table.getName()).deleteOne(key);
    }

    @Override
    public void truncate(Table table) {
        getCollection(table.getName()).drop();
    }

    @Override
    public long count(Table table) {
        return getCollection(table.getName()).countDocuments();
    }
}
